#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
MaixCAM Pro 设备上的 YOLO 对象检测摄像头应用
支持 YOLOv5/v8/v11 模型进行实时对象检测
"""

from maix import display, camera, nn, image, app
import time

class YOLOCamera:
    def __init__(self):
        """初始化摄像头、显示屏和YOLO检测器"""
        # 初始化摄像头 (640x480分辨率)
        self.cam = camera.Camera(1920, 1080)
        
        # 初始化显示屏
        self.disp = display.Display()
        
        # 初始化YOLO检测器（尝试不同版本的YOLO模型）
        self.detector = None
        self.init_yolo_detector()
        
    def init_yolo_detector(self):
        """初始化YOLO检测器，尝试加载不同版本的模型"""
        # 按优先级尝试加载模型
        yolo_models = [
            ("YOLOv11", "/root/models/yolo11n.mud", nn.YOLO11),
            ("YOLOv8", "/root/models/yolov8n.mud", nn.YOLOv8),
            ("YOLOv5", "/root/models/yolov5s.mud", nn.YOLOv5)
        ]
        
        for model_name, model_path, model_class in yolo_models:
            try:
                print(f"尝试加载{model_name}模型...")
                self.detector = model_class(model=model_path, dual_buff=True)
                print(f"{model_name}模型加载成功")
                break
            except Exception as e:
                print(f"{model_name}模型加载失败: {e}")
                continue
        
        if self.detector is None:
            print("所有YOLO模型加载失败，将只显示原始摄像头画面")
        else:
            print("支持的对象类别:")
            for i, label in enumerate(self.detector.labels):
                print(f"  {i}: {label}")
    
    def run(self):
        """主循环：捕获图像并进行对象检测"""
        print("开始摄像头对象检测...")
        print("按 Ctrl+C 停止程序")
        
        frame_count = 0
        start_time = time.time()
        
        try:
            while not app.need_exit():
                # 读取摄像头图像
                img = self.cam.read()
                
                # 如果成功加载了YOLO模型，则进行对象检测
                if self.detector:
                    try:
                        # 进行对象检测
                        objs = self.detector.detect(img, conf_th=0.5, iou_th=0.45)
                        
                        # 在图像上绘制检测结果
                        for obj in objs:
                            # 绘制边界框
                            img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED, thickness=2)
                            
                            # 显示类别标签和置信度
                            label = f'{self.detector.labels[obj.class_id]}: {obj.score:.2f}'
                            # 绘制标签背景框
                            img.draw_string(obj.x, obj.y - 20, label, color=image.COLOR_RED, scale=0.5)
                            
                            # 打印检测结果到控制台
                            print(f"检测到: {self.detector.labels[obj.class_id]}, 置信度: {obj.score:.2f}, 位置: ({obj.x}, {obj.y}, {obj.w}, {obj.h})")
                        
                        # 计算并显示帧率
                        frame_count += 1
                        if frame_count % 30 == 0:  # 每30帧计算一次帧率
                            elapsed_time = time.time() - start_time
                            fps = 30 / elapsed_time
                            print(f"帧率: {fps:.2f} FPS")
                            start_time = time.time()
                            
                    except Exception as e:
                        print(f"检测过程中出错: {e}")
                
                # 显示图像（带或不带检测结果）
                self.disp.show(img)
                
                # 添加短暂延迟以控制帧率
                time.sleep(0.01)
                
        except KeyboardInterrupt:
            print("\n用户中断程序")
        except Exception as e:
            print(f"程序运行出错: {e}")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """清理资源"""
        print("程序结束，清理资源...")
        # 清理代码（如果需要）

def main():
    """主函数"""
    print("MaixCAM Pro YOLO对象检测摄像头")
    print("=" * 40)
    
    # 创建并运行YOLO摄像头应用
    yolo_cam = YOLOCamera()
    yolo_cam.run()

if __name__ == "__main__":
    main()